'''
@Author: JintuZheng 郑晋图
@Code: Tester for mnist
'''
import torch
from torch.utils.data import random_split, DataLoader
from torch.autograd import Variable
from ours_dataloader import MnistDataset
import torch.nn.functional as F
from model import R34
import torch.nn as nn
from tqdm import trange
from matplotlib import pyplot as plt
from utils import build_loss_fig
import numpy as np

batch_szie = 100

whole_dataset = MnistDataset(root='Dataset/mnist_test/') # Make dataset
test_loader = DataLoader(whole_dataset, shuffle=True, batch_size=batch_szie)

device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') # Device set
model = R34(10) # build model

model.load_state_dict(torch.load('weights/resnet34_mnist_5.pth'))

model = model.to(device)
model.eval()

false_counter = 0
counter_len = 0

with torch.no_grad():
    for index, (img, label) in enumerate(test_loader):
        print('Testing for the index {}'.format(index))
        img, label = img.to(device), label.to(device)
        pred = model(img)
        pred = pred.detach().to('cpu').numpy()
        idx = np.argmax(pred, axis=1) # get the max index for pred class-idxs
        pro = np.amax(pred, axis=1) # pro
        label = np.argmax(label.detach().to('cpu').numpy(), axis=1)
        judege_mat = (idx == label) # make the bool matrix
        counter_len += label.shape[0] # Counetr
        false_counter += np.where(judege_mat==False)[0].shape[0] # The navigate counter

#print(counter_len)
#print(false_counter)
print('Accuracy: {}%'.format((counter_len-false_counter)*100/counter_len))        
        


